package com.smh.aiagent.agent;

import cn.hutool.core.util.StrUtil;
import com.smh.aiagent.agent.model.AgentState;
import com.smh.aiagent.enums.HttpStatus;
import com.smh.aiagent.exceptions.AiAgentException;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

/**
 * @author shiminghui
 * @date 2025/6/8 14:50
 * @description: 基础的智能体
 */

@Data
@Slf4j
public abstract class BaseAgent {

    private String name;
    private String description;

    private String SystemPrompt;
    private String nextStepPrompt;
    private String nextStepPromptReplica;
    private String stuckPrompt;

    // 智能体客户端
    private ChatClient chatClient;

    // 智能体消息列表
    private List<Message> messageList = new ArrayList<>();
    //  智能体状态
    private AgentState agentState = AgentState.IDLE;
    // 智能体最大步数
    private Integer maxStep = 10;
    // 当前步数
    private Integer currentStep = 0;
    // 重复阈值
    private Integer duplicateThreshold = 2;


    public String run(String userPrompt) {
        // 检测智能体的状态
        if (agentState != AgentState.IDLE) {
            throw new AiAgentException(HttpStatus.FAIL, "智能体状态错误");
        }
        // 检测用户提示次
        if (StrUtil.isBlank(userPrompt)) {
            throw new AiAgentException(HttpStatus.FAIL, "用户提示不能为空");
        }
        // 修改智能体状态
        agentState = AgentState.RUNNING;
        // 添加用户提示到智能体消息列表
        messageList.add(new UserMessage(userPrompt));
        // 存储运行结果
        List<String> results = new ArrayList<>();

        try {
            while (currentStep < maxStep && agentState == AgentState.RUNNING) {
                // 将当前步数加一
                currentStep++;
                // 打个日志记录一下
                log.info("当前步数:{}, 最大步数:{}", currentStep, maxStep);
                // 获取下一步
                String result = getNextStep();
                log.info("当前步数:{},运行完成,运行结果为:{}", currentStep, result);
                // 添加到结果列表中
                results.add(result);

                if (isStuck()) {
                    handlerStuck();
                } else {
                    nextStepPrompt = nextStepPromptReplica;
                }
            }

            if (currentStep >= maxStep) {
                agentState = AgentState.FINISHED;
            }
            return String.join("\n", results);
        } catch (Exception e) {
            log.error("[{}]执行错误,错误信息为:{}", name, e.getMessage());
            return "执行错误,错误信息为:" + e.getMessage();
        } finally {
            cleanup();
        }
    }

    /**
     * 运行代理（流式输出）
     *
     * @param userPrompt 用户提示词
     * @return SseEmitter实例
     */
    public SseEmitter runStream(String userPrompt) {
        // 创建SseEmitter，设置较长的超时时间
        SseEmitter emitter = new SseEmitter(300000L); // 5分钟超时

        // 使用线程异步处理，避免阻塞主线程
        CompletableFuture.runAsync(() -> {
            try {
                if (this.agentState != AgentState.IDLE) {
                    emitter.send("错误：无法从状态运行代理: " + this.agentState);
                    emitter.complete();
                    return;
                }
                if (StrUtil.isBlank(userPrompt)) {
                    emitter.send("错误：不能使用空提示词运行代理");
                    emitter.complete();
                    return;
                }

                // 更改状态
                agentState = AgentState.RUNNING;
                // 记录消息上下文
                messageList.add(new UserMessage(userPrompt));

                try {
                    for (int i = 0; i < maxStep && agentState != AgentState.FINISHED; i++) {
                        int stepNumber = i + 1;
                        currentStep = stepNumber;
                        log.info("Executing step " + stepNumber + "/" + maxStep);

                        // 单步执行
                        String stepResult = getNextStep();
                        String result = "Step " + stepNumber + ": " + stepResult;

                        // 发送每一步的结果
                        emitter.send(result);
                    }
                    // 检查是否超出步骤限制
                    if (currentStep >= maxStep) {
                        agentState = AgentState.FINISHED;
                        emitter.send("执行结束: 达到最大步骤 (" + maxStep + ")");
                    }
                    // 正常完成
                    emitter.complete();
                } catch (Exception e) {
                    agentState = AgentState.ERROR;
                    log.error("执行智能体失败", e);
                    try {
                        emitter.send("执行错误: " + e.getMessage());
                        emitter.complete();
                    } catch (Exception ex) {
                        emitter.completeWithError(ex);
                    }
                } finally {
                    // 清理资源
                    this.cleanup();
                }
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
        });

        // 设置超时和完成回调
        emitter.onTimeout(() -> {
            this.agentState = AgentState.ERROR;
            this.cleanup();
            log.warn("SSE connection timed out");
        });

        emitter.onCompletion(() -> {
            if (this.agentState == AgentState.RUNNING) {
                this.agentState = AgentState.FINISHED;
            }
            this.cleanup();
            log.info("SSE connection completed");
        });

        return emitter;
    }


    protected abstract String getNextStep();

    protected void cleanup() {
        agentState = AgentState.IDLE;
        messageList.clear();
        currentStep = 0;
    }

    /**
     * 判断是否卡住
     *
     * @return
     */
    protected boolean isStuck() {
        if (getMessageList().size() < 2) {
            return false;
        }

        Message last = getMessageList().getLast();
        if (last == null && last.getText().isEmpty()) {
            return false;
        }

        int duplicatedStep = 0;

        if (last.getMessageType().equals(MessageType.TOOL)) {
            ToolResponseMessage toolResponseMessage = (ToolResponseMessage) last;
            if (toolResponseMessage.getResponses().isEmpty()) {
                return false;
            }
            List<ToolResponseMessage.ToolResponse> responses = toolResponseMessage.getResponses();
            for (int i = getMessageList().size() - 2; i >= 0; i--) {
                Message previous = getMessageList().get(i);
                if (!previous.getMessageType().equals(MessageType.TOOL)) {
                    continue;
                }
                ToolResponseMessage previousToolResponseMessage = (ToolResponseMessage) previous;
                if (previousToolResponseMessage.getResponses().size() != responses.size()) {
                    return false;
                }
                for (int j = 0; j < responses.size(); j++) {
                    if (!responses.get(j).responseData().equals(previousToolResponseMessage.getResponses().get(j).responseData())) {
                        return false;
                    }
                    duplicatedStep++;
                }
            }

        } else {
            for (int i = getMessageList().size() - 2; i >= 0; i--) {
                Message previous = getMessageList().get(i);
                if (previous.getMetadata().equals(last.getMetadata()) && previous.getText().equals(last.getText())) {
                    duplicatedStep++;
                }
            }
        }

        return duplicatedStep >= duplicateThreshold;
    }

    protected void handlerStuck() {
        this.nextStepPrompt = stuckPrompt + "\n" + nextStepPrompt;
        log.info("检测到重复操作, 正在纠正.");
    }

}
